/*
  File:    WENO.cpp
  Purpose: weighted ENO (essentiall non-oscillatory) schemes for discretizing 
           |Speed * grad(u)|   

  Contact: Paul Macklin
           pmacklin@math.uci.edu
		  http://math.uci.edu/~pmacklin
*/

#include <cmath>
#include <cstdlib>
#include <cstdio>
#include <iostream>

using namespace std;

#include "../ImprovedMath.h"
#include "../Matrix.h"
#include "WENO.h" 

// this one does mSpeedIn | grad( mIn ) |

Matrix WENO3(Matrix &mSpeedIn, Matrix &mIn)
{
 int i=0, j=0;
 Matrix mMatrixOut = zeros(mIn.TellRows(), mIn.TellCols()); 
 
 for(i=0+3 ; i < mMatrixOut.TellRows()-3 ; i++)
 {
  for(j=0+3 ; j < mMatrixOut.TellCols()-3 ; j++)
  {
   // First, we'll do the x-derivative
  
   double dTemp1 = CentDiffXX( mIn, i-1, j );
   double dTemp2 = CentDiffXX( mIn, i, j);
   
   double dTemp3 = ForwardDiffX( mIn, i-2, j);
   double dTemp4 = ForwardDiffX( mIn, i-1, j);
   double dTemp5 = ForwardDiffX( mIn, i,   j);
   double dTemp6 = ForwardDiffX( mIn, i+1, j);
   double dTemp7 = ForwardDiffX( mIn, i+2, j);
      
   double dR = ( _EPS + square( dTemp1 ) ) / (_EPS + square(dTemp2) );
   double dOmegaMinus = 1.0 / (1.0 + 2.0*square(dR) );
   double dLeftXDeriv = ( ( dTemp4 + dTemp5 ) - 
              dOmegaMinus*( dTemp3 - 2.0*dTemp4 + dTemp5)
                        )/
                        (2.0*dx);
                        
   dTemp1 = CentDiffXX( mIn, i+1, j );
   
   dR = ( _EPS + square( dTemp1 ) ) / (_EPS + square(dTemp2) );
   double dOmegaPlus = 1.0 / (1.0 + 2.0*square(dR) );                         
   double dRightXDeriv = ( ( dTemp4 + dTemp5 ) - 
                dOmegaPlus*( dTemp6 - 2.0*dTemp5 + dTemp4)
                         )/
                         (2.0*dx);
  
   dTemp1 = CentDiffYY( mIn, i, j-1);
   dTemp2 = CentDiffYY( mIn, i,   j);
   
   dTemp3 = ForwardDiffY( mIn, i, j-2);
   dTemp4 = ForwardDiffY( mIn, i, j-1);
   dTemp5 = ForwardDiffY( mIn, i,   j);
   dTemp6 = ForwardDiffY( mIn, i, j+1); 
      
   dR = ( _EPS + square( dTemp1 ) ) / (_EPS + square(dTemp2) );
   dOmegaMinus = 1.0 / (1.0 + 2.0*square(dR) );
   double dLeftYDeriv = ( ( dTemp4 + dTemp5 ) - 
              dOmegaMinus*( dTemp3 - 2.0*dTemp4 + dTemp5)
                        )/
                        (2.0*dy);
                        
   dTemp1 = CentDiffYY( mIn, i, j+1);
   
   dR = ( _EPS + square( dTemp1 ) ) / (_EPS + square(dTemp2 ) );
   dOmegaPlus = 1.0 / (1.0 + 2.0*square(dR) );                         
   double dRightYDeriv = ( ( dTemp4 + dTemp5 ) - 
                dOmegaPlus*( dTemp6 - 2.0*dTemp5 + dTemp4)
                         )/
                         (2.0*dy);                         
                         
   double dPlusGrad = sqrt(  1*square( PlusFunction(dLeftXDeriv) )
                           + 1*square( MinusFunction(dRightXDeriv) )
                           + 1*square( PlusFunction(dLeftYDeriv) )
                           + 1*square( MinusFunction(dRightYDeriv) ) );
   double dMinusGrad = sqrt(  1*square( MinusFunction(dLeftXDeriv) )
                            + 1*square( PlusFunction(dRightXDeriv) )
                            + 1*square( MinusFunction(dLeftYDeriv) )
                            + 1*square( PlusFunction(dRightYDeriv) ) );                           
   dTemp1 = *mSpeedIn(i,j);                            
   *mMatrixOut(i,j) = PlusFunction( dTemp1 ) * dPlusGrad + 
                        1*MinusFunction( dTemp1 ) * dMinusGrad;
  }
 }
 
 
 return mMatrixOut;
}

double WENOweight( double dA, double dB, double dC, double dD)                 
{
 double dIS0 = IS0( dA, dB, dC, dD);
 double dIS1 = IS1( dA, dB, dC, dD);
 double dIS2 = IS2( dA, dB, dC, dD);
 
 double dAlpha0 = 1.0 / (_EPS+square(dIS0));
 double dAlpha1 = 6.0 / (_EPS+square(dIS1));
 double dAlpha2 = 3.0 / (_EPS+square(dIS2));
 
 double dSum = dAlpha0 + dAlpha1 + dAlpha2 + _EPS;
 double dOmega0 = (dAlpha0+_EPS) / dSum;
 double dOmega2 = (dAlpha2+_EPS) / dSum;
 
 double dOut = (1.0/3.0)*dOmega0*(dA-2.0*dB+dC);
 dOut = dOut + (1.0/6.0)*(dOmega2-.5)*(dB-2.0*dC+dD);
 return dOut;
}

// this one does mSpeedIn | grad( mIn ) |

Matrix WENO5(Matrix &mSpeedIn, Matrix &mIn)
{
 int i=0, j=0;
 Matrix mMatrixOut = zeros(mIn.TellRows(), mIn.TellCols()); 

 for(i=0+3 ; i < mMatrixOut.TellRows()-3 ; i++)
 {
  for(j=0+3 ; j < mMatrixOut.TellCols()-3 ; j++)
  {
   // First, we'll do the x-derivative
   
   double dTemp1, dTemp2, dTemp3, dTemp4, dTemp5, 
          dTemp6, dTemp7, dTemp8, dTemp9;
  
   dTemp1 = ForwardDiffX( mIn , i-2, j);       
   dTemp2 = ForwardDiffX( mIn , i-1, j);        
   dTemp3 = ForwardDiffX( mIn , i  , j);
   dTemp4 = ForwardDiffX( mIn , i+1, j);
   
   dTemp5 = CentDiffXX( mIn, i-2, j);
   dTemp6 = CentDiffXX( mIn, i-1, j);
   dTemp7 = CentDiffXX( mIn, i  , j);
   dTemp8 = CentDiffXX( mIn, i+1, j);
   dTemp9 = CentDiffXX( mIn, i+2, j);
   
   double dLeftXDeriv = ( -dTemp1 + 7.0 * (dTemp2 + dTemp3) -dTemp4)/(12.0*dx);
   double dRightXDeriv = dLeftXDeriv;
   dLeftXDeriv = dLeftXDeriv - WENOweight( dTemp5/dx, dTemp6/dx,
                                           dTemp7/dx, dTemp8/dx);
   dRightXDeriv = dRightXDeriv + WENOweight( dTemp9/dx, dTemp8/dx,
                                             dTemp7/dx, dTemp6/dx);                           

   dTemp1 = ForwardDiffY( mIn , i , j-2);       
   dTemp2 = ForwardDiffY( mIn , i , j-1);        
   dTemp3 = ForwardDiffY( mIn , i , j  );
   dTemp4 = ForwardDiffY( mIn , i , j+1);
   
   dTemp5 = CentDiffYY( mIn, i , j-2);
   dTemp6 = CentDiffYY( mIn, i , j-1);
   dTemp7 = CentDiffYY( mIn, i , j  );
   dTemp8 = CentDiffYY( mIn, i , j+1);
   dTemp9 = CentDiffYY( mIn, i , j+2);
   
   double dLeftYDeriv = ( -dTemp1 + 7.0 * (dTemp2 + dTemp3) -dTemp4)/(12.0*dy);
   double dRightYDeriv = dLeftYDeriv;
   dLeftYDeriv = dLeftYDeriv - WENOweight( dTemp5/dy, dTemp6/dy,
                                           dTemp7/dy, dTemp8/dy);
   dRightYDeriv = dRightYDeriv + WENOweight( dTemp9/dy, dTemp8/dy,
                                                    dTemp7/dy, dTemp6/dy); 
                                                                                                      
   dTemp1 = *mSpeedIn(i,j);   
   if( dTemp1 > _EPS ) 
   {
    double dPlusGrad = sqrt(  square( PlusFunction(dLeftXDeriv) )
                            + square( MinusFunction(dRightXDeriv) )
                            + square( PlusFunction(dLeftYDeriv) )
                            + square( MinusFunction(dRightYDeriv) ) + _EPS );
    *mMatrixOut(i,j) = dTemp1 * dPlusGrad;
   }
   else
   {
    double dMinusGrad = sqrt(  square( MinusFunction(dLeftXDeriv) )
                             + square( PlusFunction(dRightXDeriv) )
                             + square( MinusFunction(dLeftYDeriv) )
                             + square( PlusFunction(dRightYDeriv) ) + _EPS );
    *mMatrixOut(i,j) = dTemp1 * dMinusGrad; 
   }

  }
 }
 
 return mMatrixOut;
}

Matrix WENO5banded(Matrix &mSpeedIn, Matrix &mIn, double BandSize )
{
 int i=0, j=0;
 Matrix mMatrixOut = zeros(mIn.TellRows(), mIn.TellCols()); 

 for(i=0+3 ; i < mMatrixOut.TellRows()-3 ; i++)
 {
  for(j=0+3 ; j < mMatrixOut.TellCols()-3 ; j++)
  {
   // First, we'll do the x-derivative
   
   if( *mIn(i,j) < BandSize )
   {
    double dTemp1, dTemp2, dTemp3, dTemp4, dTemp5, 
           dTemp6, dTemp7, dTemp8, dTemp9;
   
    dTemp1 = ForwardDiffX( mIn , i-2, j);       
    dTemp2 = ForwardDiffX( mIn , i-1, j);        
    dTemp3 = ForwardDiffX( mIn , i  , j);
    dTemp4 = ForwardDiffX( mIn , i+1, j);
    
    dTemp5 = CentDiffXX( mIn, i-2, j);
    dTemp6 = CentDiffXX( mIn, i-1, j);
    dTemp7 = CentDiffXX( mIn, i  , j);
    dTemp8 = CentDiffXX( mIn, i+1, j);
    dTemp9 = CentDiffXX( mIn, i+2, j);
    
    double dLeftXDeriv = ( -dTemp1 + 7.0 * (dTemp2 + dTemp3) -dTemp4)/(12.0*dx);
    double dRightXDeriv = dLeftXDeriv;
    dLeftXDeriv = dLeftXDeriv - WENOweight( dTemp5/dx, dTemp6/dx,
                                            dTemp7/dx, dTemp8/dx);
    dRightXDeriv = dRightXDeriv + WENOweight( dTemp9/dx, dTemp8/dx,
                                              dTemp7/dx, dTemp6/dx);                            
         
    dTemp1 = ForwardDiffY( mIn , i , j-2);       
    dTemp2 = ForwardDiffY( mIn , i , j-1);        
    dTemp3 = ForwardDiffY( mIn , i , j  );
    dTemp4 = ForwardDiffY( mIn , i , j+1);
   
    dTemp5 = CentDiffYY( mIn, i , j-2);
    dTemp6 = CentDiffYY( mIn, i , j-1);
    dTemp7 = CentDiffYY( mIn, i , j  );
    dTemp8 = CentDiffYY( mIn, i , j+1);
    dTemp9 = CentDiffYY( mIn, i , j+2);
    
    double dLeftYDeriv = ( -dTemp1 + 7.0 * (dTemp2 + dTemp3) -dTemp4)/(12.0*dy);
    double dRightYDeriv = dLeftYDeriv;
    dLeftYDeriv = dLeftYDeriv - WENOweight( dTemp5/dy, dTemp6/dy,
                                            dTemp7/dy, dTemp8/dy);
    dRightYDeriv = dRightYDeriv + WENOweight( dTemp9/dy, dTemp8/dy,
                                              dTemp7/dy, dTemp6/dy); 
                                                                                                       
    dTemp1 = *mSpeedIn(i,j);   
    if( dTemp1 > _EPS ) 
    {
     double dPlusGrad = sqrt(  square( PlusFunction(dLeftXDeriv) )
                             + square( MinusFunction(dRightXDeriv) )
                             + square( PlusFunction(dLeftYDeriv) )
                             + square( MinusFunction(dRightYDeriv) ) + _EPS );
     *mMatrixOut(i,j) = dTemp1 * dPlusGrad;
    }
    else
    {
     double dMinusGrad = sqrt(  square( MinusFunction(dLeftXDeriv) )
                              + square( PlusFunction(dRightXDeriv) )
                              + square( MinusFunction(dLeftYDeriv) )
                              + square( PlusFunction(dRightYDeriv) ) + _EPS );
     *mMatrixOut(i,j) = dTemp1 * dMinusGrad; 
    }
   }
    
  }
 }
 
 return mMatrixOut;
}

// This does (u,v).grad(f)

Matrix WENO5VectorSpeed(Matrix &mSpeedInX, Matrix &mSpeedInY, Matrix &mIn)
{
 int i=0, j=0;
 Matrix mMatrixOut = zeros(mIn.TellRows(), mIn.TellCols()); 
 
 for(i=0+3 ; i < mMatrixOut.TellRows()-3 ; i++)
 {
  for(j=0+3 ; j < mMatrixOut.TellCols()-3 ; j++)
  {
   double dXTerm = *mSpeedInX(i,j);
   double dYTerm = *mSpeedInY(i,j);

   // First, we'll do the x-derivative
   
   double dTemp1, dTemp2, dTemp3, dTemp4, dTemp5, 
          dTemp6, dTemp7, dTemp8, dTemp9;
  
   dTemp1 = ForwardDiffX( mIn , i-2, j);       
   dTemp2 = ForwardDiffX( mIn , i-1, j);        
   dTemp3 = ForwardDiffX( mIn , i  , j);
   dTemp4 = ForwardDiffX( mIn , i+1, j);
   
   dTemp5 = CentDiffXX( mIn, i-2, j);
   dTemp6 = CentDiffXX( mIn, i-1, j);
   dTemp7 = CentDiffXX( mIn, i  , j);
   dTemp8 = CentDiffXX( mIn, i+1, j);
   dTemp9 = CentDiffXX( mIn, i+2, j);
   
   double dLeftXDeriv = ( -dTemp1 + 7.0 * (dTemp2 + dTemp3) -dTemp4)/(12.0*dx);
   double dRightXDeriv = dLeftXDeriv;

   if( dYTerm >= _EPS )
   {
    dLeftXDeriv = dLeftXDeriv - WENOweight( dTemp5/dx, dTemp6/dx,
                                            dTemp7/dx, dTemp8/dx);
    dXTerm *= dLeftXDeriv;
   }
   else
   {
    dRightXDeriv = dRightXDeriv + WENOweight( dTemp9/dx, dTemp8/dx,
                                              dTemp7/dx, dTemp6/dx);
    dXTerm *= dRightXDeriv;
   }

   dTemp1 = ForwardDiffY( mIn , i , j-2);       
   dTemp2 = ForwardDiffY( mIn , i , j-1);        
   dTemp3 = ForwardDiffY( mIn , i , j  );
   dTemp4 = ForwardDiffY( mIn , i , j+1);
   
   dTemp5 = CentDiffYY( mIn, i , j-2);
   dTemp6 = CentDiffYY( mIn, i , j-1);
   dTemp7 = CentDiffYY( mIn, i , j  );
   dTemp8 = CentDiffYY( mIn, i , j+1);
   dTemp9 = CentDiffYY( mIn, i , j+2);
   
   double dLeftYDeriv = ( -dTemp1 + 7.0 * (dTemp2 + dTemp3) -dTemp4)/(12.0*dy);
   double dRightYDeriv = dLeftYDeriv;

   if( dYTerm >= _EPS )
   {
    dLeftYDeriv = dLeftYDeriv - WENOweight( dTemp5/dy, dTemp6/dy,
                                            dTemp7/dy, dTemp8/dy);
    dYTerm *= dLeftYDeriv;
   }
   else
   {
    dRightYDeriv = dRightYDeriv + WENOweight( dTemp9/dy, dTemp8/dy,
                                              dTemp7/dy, dTemp6/dy);
    dYTerm *= dRightYDeriv;
   }
              
   *mMatrixOut(i,j) = dXTerm + dYTerm;
  }
 }
 
 return mMatrixOut;
}
